import os
import json
from fastapi import HTTPException
from sqlalchemy.orm import Session
import time

from app.ml.predict import predict
from app.db_models.prediction_log import PredictionLog

# 模型和结果保存路径
MODEL_DIR = "models"
RESULTS_DIR = "results"
PREDICTIONS_DIR = "predictions"

def get_available_models():
    """获取所有可用的模型"""
    models = []
    if os.path.exists(MODEL_DIR):
        for file in os.listdir(MODEL_DIR):
            if file.endswith("_model.h5"):
                model_type = file.split("_")[0]  # 从文件名中提取模型类型
                models.append({
                    "name": model_type,
                    "path": os.path.join(MODEL_DIR, file)
                })
    return models

def run_prediction(model_type, data_path):
    """运行预测并返回结果"""
    model_path = os.path.join(MODEL_DIR, f"{model_type}_model.h5")
    
    # 检查模型是否存在
    if not os.path.exists(model_path):
        raise HTTPException(
            status_code=404,
            detail=f"模型 {model_type} 不存在，请先训练模型"
        )
    
    # 调用预测函数
    result = predict(model_path, data_path)
    return result

def get_latest_prediction_result():
    """获取最新的预测结果"""
    latest_file = None
    latest_time = 0
    
    if os.path.exists(PREDICTIONS_DIR):
        for file in os.listdir(PREDICTIONS_DIR):
            if file.startswith("prediction_result_") and file.endswith(".json"):
                file_path = os.path.join(PREDICTIONS_DIR, file)
                file_time = os.path.getmtime(file_path)
                if file_time > latest_time:
                    latest_time = file_time
                    latest_file = file_path
    
    if latest_file:
        with open(latest_file, 'r') as f:
            return json.load(f)
    
    return None

def save_prediction_to_db(results, user_id, db: Session):
    """将预测结果保存到数据库"""
    try:
        # 从结果中提取需要的信息
        model_type = results["model_info"]["model_type"]
        dataset_path = results["data_info"]["data_path"]
        dataset_name = os.path.basename(dataset_path)
        
        evaluation = results["prediction_info"]["evaluation"]
        accuracy = evaluation["accuracy"] * 100  # 转为百分比
        precision_score = evaluation["precision"] * 100
        recall_score = evaluation["recall"] * 100
        f1_score = evaluation["f1_score"] * 100
        loss = evaluation.get("loss", 0)
        
        sample_count = results["data_info"]["n_samples"]
        feature_count = results["data_info"]["n_features"]
        class_distribution = json.dumps(results["prediction_info"]["class_distribution"])
        
        # 计算类别数
        class_count = len(results["prediction_info"]["class_distribution"])
        
        prediction_time = results["prediction_info"]["prediction_time"]
        
        # 创建日志文件路径
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        result_file_path = os.path.join(PREDICTIONS_DIR, f"prediction_result_{timestamp}.json")
        
        # 创建新记录
        log = PredictionLog(
            user_id=user_id,
            model_type=model_type,
            dataset_name=dataset_name,
            dataset_path=dataset_path,
            accuracy=accuracy,
            precision_score=precision_score,
            recall_score=recall_score,
            f1_score=f1_score,
            loss=loss,
            sample_count=sample_count,
            feature_count=feature_count,
            class_count=class_count,
            class_distribution=class_distribution,
            prediction_time=prediction_time,
            result_file_path=result_file_path
        )
        
        # 保存到数据库
        db.add(log)
        db.commit()
        db.refresh(log)
        
        return log.id
    except Exception as e:
        db.rollback()
        print(f"保存预测日志失败: {str(e)}")
        return None